'''
========================================
-->    MISC JUNK
========================================
'''
def objExists( inObjName ):
    '''    
    Takes an object name and returns True if it exists in the scene.
    False otherwise.
    
    @param inObjName: String.
    '''
    selList = OpenMaya.MSelectionList()
    try:
        selList.add( inObjName )
    except:
        return False
        
    if not selList.isEmpty():
        return True
    else:
        return False
    
def getMObject( inObjName ):
    '''
    Convert object string name into MObject.
    
    @param inObjName: String.
    @return: MObject.
    '''
    selList = OpenMaya.MSelectionList()
    if inObjName[-1:] is not '*':
        try:
            selList.add( inObjName )
            obj = OpenMaya.MObject()
            selList.getDependNode( 0, obj )
            return obj
        except:
            pass

def createNode( inNodeName=None, inNodeType=None, inNodeParent=None ):
    """
    Creates a node of the given type.
    
    @param: inNodeName. String. Name of the new node.
    @param inNodeType: String. Type of node to create.
    @param inNodeParent: String or MObject. Name of node's parent.
    @return: The new node.
    """
    if not isinstance( inNodeType, str ):
        OpenMaya.MGlobal.displayError( 'createNode argument inNodeType must be a string.' )
        return
        
    if inNodeParent:
        if isinstance( inNodeParent, str ):
            if objExists( inNodeParent ):
                parentObj = getMObject( inNodeParent )
            
            else:
                OpenMaya.MGlobal.displayError( 'Object {0} doesn\t exist'.format( inNodeParent ) )
                return

        elif isinstance( inNodeParent, OpenMaya.MObject ):
            parentObj = inNodeParent
            
        else:
            OpenMaya.MGlobal.displayError( 'createNode argument inNodeParent must be a String or MObject.' )
            return
        
    
    mdgMod = OpenMaya.MDGModifier()
    
    if inNodeParent:
        dagMod = OpenMaya.MDagModifier()
        node = dagMod.createNode( inNodeType, parentObj )
        dagMod.doIt()
        
    else:
        node = mdgMod.createNode( inNodeType )  
    
    if not inNodeName:
        inNodeName = inNodeType+'1'
    
    mdgMod.renameNode( node, inNodeName )
    mdgMod.doIt()
    
    return node

def getPlugByName( inObj, inPlugName ):
    """
    Gets a node's plug as an MPlug.
    
    @param inObj: MObject. Node to get plug from.
    @param inPlugName: String. Name of plug to get from node.
    @return: MPlug.
    @return: None.
    """
    depFn = OpenMaya.MFnDependencyNode( inObj )
    
    try:
        plug = depFn.findPlug( inPlugName )
        return plug
    
    except:
        return None
    
def setPlugValue( inPlug, inValue ):
    """
    Sets the given plug's value to the passed in value.

    @param inPlug: _MPlug_. The node plug.
    @param inValue: _Type_. Any value of any data type.
    """
    plugAttribute = inPlug.attribute()
    apiType = plugAttribute.apiType()
    #print plugAttribute.apiTypeStr()
    
    # Float Groups - rotate, translate, scale
    if apiType in [ OpenMaya.MFn.kAttribute3Double, OpenMaya.MFn.kAttribute3Float ]:
        result = []
        
        if inPlug.isCompound():
            if isinstance( inValue, list ):
                for c in xrange( inPlug.numChildren() ):
                    result.append( setPlugValue( inPlug.child( c ), inValue[ c ] ) )
                    
                return result
            
            elif type( inValue ) == OpenMaya.MEulerRotation:
                setPlugValue( inPlug.child( 0 ), inValue.x )
                setPlugValue( inPlug.child( 1 ), inValue.y )
                setPlugValue( inPlug.child( 2 ), inValue.z )
                
            else:
                OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type list.'.format( inPlug.info(), inValue, type( inValue ) ) )
    
    # Distance
    elif apiType in [ OpenMaya.MFn.kDoubleLinearAttribute, OpenMaya.MFn.kFloatLinearAttribute ]:
        if isinstance( inValue, float ):
            value =  OpenMaya.MDistance( inValue, OpenMaya.MDistance.kCentimeters )
            inPlug.setMDistance( value )
            
        else:
            OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type float.'.format( inPlug.info(), inValue, type( inValue ) ) )
    
    # Angle
    elif apiType in [ OpenMaya.MFn.kDoubleAngleAttribute, OpenMaya.MFn.kFloatAngleAttribute ]:
        if isinstance( inValue, float ):
            value = OpenMaya.MAngle( inValue, OpenMaya.MAngle.kDegrees )
            inPlug.setMAngle( value )
            
        else:
            OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type float.'.format( inPlug.info(), inValue, type( inValue ) ) )
            
    # Typed - matrix WE DON'T HANDLE THIS CASE YET!!!!!!!!!
    elif apiType == OpenMaya.MFn.kTypedAttribute:
        pType = OpenMaya.MFnTypedAttribute( plugAttribute ).attrType()
        
        if pType == OpenMaya.MFnData.kMatrix:
            if isinstance( inValue, OpenMaya.MPlug ):
                pass
            
            else:
                plugNode = inPlug.node()
                MFnTrans = OpenMaya.MFnTransform( plugNode )
                sourceMatrix = OpenMaya.MTransformationMatrix( inValue )#.asMatrix()
                MFnTrans.set( sourceMatrix )
        
        # String
        elif pType == OpenMaya.MFnData.kString:
            value = inValue
            inPlug.setString( value )
    
    # MATRIX
    elif apiType == OpenMaya.MFn.kMatrixAttribute:
        if isinstance( inValue, OpenMaya.MPlug ):
            # inValue must be a MPlug!
            sourceValueAsMObject = OpenMaya.MFnMatrixData( inValue.asMObject() ).object()
            inPlug.setMObject( sourceValueAsMObject )
            
        else:
            OpenMaya.MGlobal.displayError( 'Value object is not an MPlug. To set a MMatrix value, both passed in variables must be MPlugs.' )
    
    # Numbers
    elif apiType == OpenMaya.MFn.kNumericAttribute:
        pType = OpenMaya.MFnNumericAttribute( plugAttribute ).unitType()
        
        if pType == OpenMaya.MFnNumericData.kBoolean:
            if isinstance( inValue, bool ):
                inPlug.setBool( inValue )
                
            else:
                OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type bool.'.format( inPlug.info(), inValue, type( inValue ) ) )
                
        elif pType in [ OpenMaya.MFnNumericData.kShort, OpenMaya.MFnNumericData.kInt, OpenMaya.MFnNumericData.kLong, OpenMaya.MFnNumericData.kByte ]:
            if isinstance( inValue, int ):
                inPlug.setInt( inValue )
            
            else:
                OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type int.'.format( inPlug.info(), inValue, type( inValue ) ) )
        
        elif pType in [ OpenMaya.MFnNumericData.kFloat, OpenMaya.MFnNumericData.kDouble, OpenMaya.MFnNumericData.kAddr ]:
            if isinstance( inValue, float ):
                inPlug.setDouble( inValue )
           
            else:
                OpenMaya.MGlobal.displayError( '{0} :: Passed in value ( {1} ) is {2}. Needs to be type float.'.format( inPlug.info(), inValue, type( inValue ) ) )
                
    # Enums
    elif apiType == OpenMaya.MFn.kEnumAttribute:
        inPlug.setInt( inValue )
        
def connectNodes( inParentObj, inParentPlug, inChildObj, inChildPlug ):
    '''
    @param inParentObj: String. Name of parent node.
    @param inParentPlug: String. Name of plug on parent node.
    @param inChildObj: String. Name of child node.
    @param inChildPlug: String. Name of plug on child node.
    '''
    parentPlug = getPlugByName( inParentObj, inParentPlug )
    childPlug = getPlugByName( inChildObj, inChildPlug )
    
    MDGMod = OpenMaya.MDGModifier()
    MDGMod.connect( parentPlug, childPlug )
    MDGMod.doIt()
        
def getPlugConnections( plug, dest, src ):
    '''
    @param plug: MPlug.
    @param dest: Bool. True to check plug as a destination.
    @param src: Bool. True to check the plug as a source.
    @return: MPlugArray.
    '''
    plug_array = OpenMaya.MPlugArray()
    
    if plug.connectedTo( plug_array, dest, src ):
        return plug_array
    else:
        return None

        
'''
========================================
-->    IMPORTS
========================================
'''
import maya.OpenMaya as OpenMaya       


'''
============================================================
--->>   META NODE CLASS
============================================================
'''
class BaseMeta( object ):
    def __init__( self, node_name=None, node_type='BaseMeta' ):
        '''
        @param node_name: String. Name of the node.
        '''
        self.version = 1.0
        self.type = node_type
        
        # Get or make the node.
        if objExists( node_name ):
            obj = getMObject( node_name )
            if obj.apiType() == OpenMaya.MFn.kAffect:
                self.mobject = obj
                self.name = node_name
        else:
            self.mobject = createNode( inNodeName=node_name, inNodeType='network' )
            self.name = node_name
            
            # Add the nodes attributes.
            self.initAttributes()
            
    def initAttributes( self ):
        '''
        Default attributes for this node.
        '''
        dep_fn = OpenMaya.MFnDependencyNode()
        dep_fn.setObject( self.mobject )
        
        # Version.
        version_num_attr = OpenMaya.MFnNumericAttribute()
        version_obj = version_num_attr.create( 'version', 'version', OpenMaya.MFnNumericData.kFloat )
        version_num_attr.setDefault( self.version )
        version_num_attr.setHidden( False )
        version_num_attr.setDefault( self.version )
        dep_fn.addAttribute( version_obj )
        
        # Name.
        name_str_attr = OpenMaya.MFnTypedAttribute()
        name_obj = name_str_attr.create( 'name', 'name', OpenMaya.MFnData.kString )
        name_str_attr.setHidden( False )
        dep_fn.addAttribute( name_obj )
        self.setAttributeValue( 'name', self.name )
        
        # Type.
        type_str_attr = OpenMaya.MFnTypedAttribute()
        type_obj = type_str_attr.create( 'type', 'type', OpenMaya.MFnData.kString )
        type_str_attr.setHidden( False )
        dep_fn.addAttribute( type_obj )
        self.setAttributeValue( 'type', self.type )
                
    def setAttributeValue( self, attr_name, new_value ):
        '''
        Sets the attribute to the desired value.
        
        @param attr_name: String.
        @param new_value: Type is MPlug dependent.
        '''
        plug = getPlugByName( self.mobject, attr_name )
        setPlugValue( plug, new_value )
        

class RigMeta( BaseMeta ):
    def __init__( self, node_name=None ):
        super( RigMeta, self ).__init__( node_name=node_name, node_type='RigMeta' )
        
    def initAttributes( self ):
        '''
        Default attributes for this node.
        '''
        super( RigMeta, self ).initAttributes()
        dep_fn = OpenMaya.MFnDependencyNode()
        dep_fn.setObject( self.mobject )
        
        # Meta children.
        children_str_attr = OpenMaya.MFnTypedAttribute()
        children_obj = children_str_attr.create( 'metaChildren', 'metaChildren', OpenMaya.MFnData.kString )
        children_str_attr.setArray( True )
        dep_fn.addAttribute( children_obj )
        
    def addChildMeta( self, source_plug ):
        '''
        This is an array plug so we need to handle connecting to it differently.
        We use logical index to create a new element. We also need to set some 
        data on the newly created element otherwise Maya won't physically create
        the element.
        
        @param source_plug: MPlug.
        '''
        dep_fn = OpenMaya.MFnDependencyNode( self.mobject )
        children_plugs = dep_fn.findPlug( 'metaChildren' )
        children_elements = children_plugs.numElements()
        
        # Check if the source plug is already connected.
        for plug_idx in xrange( children_elements ):
            plug = children_plugs[ plug_idx ]            
            connected_plugs = getPlugConnections( plug, True, False )
            for connected_idx in xrange( connected_plugs.length() ):
                connected_plug = connected_plugs[ connected_idx ]
                if connected_plug.name() == source_plug.name():
                    # Already connected. Get out of the function.
                    OpenMaya.MGlobal.displayInfo( '{0} is already connected to {1}'.format( source_plug.name(), plug.name() ) ) 
                    return
        
        # Make the new element plug.
        destination_plug = children_plugs.elementByLogicalIndex( children_elements )
        
        # Do the connection.
        MDGMod = OpenMaya.MDGModifier()
        MDGMod.connect( source_plug, destination_plug )
        MDGMod.doIt()
        
    def getMetaNodes( self ):
        '''
        Returns a list of all meta nodes associated with this rig.
        '''
        results = []
        
        dep_fn = OpenMaya.MFnDependencyNode( self.mobject )
        children_plugs = dep_fn.findPlug( 'metaChildren' )
        
        for plug_idx in xrange( children_plugs.numElements() ):
            plug = children_plugs[ plug_idx ]
            connected_plugs = getPlugConnections( plug, True, False )
            
            for connected_idx in xrange( connected_plugs.length() ):
                connected_plug = connected_plugs[ connected_idx ]
                node = connected_plug.node()
                type_plug = getPlugByName( node, 'type' )
                
                type_name = type_plug.asString()
                node_name = connected_plug.asString()                
                new_class = globals()[ type_name ]
                results.append( new_class( node_name=node_name ) )
        
        return results
        

class ArmMeta( BaseMeta ):
    def __init__( self, node_name=None ):
        super( ArmMeta, self ).__init__( node_name=node_name, node_type='ArmMeta' )
        
    def initAttributes( self ):
        '''
        Default attributes for this node.
        '''
        super( ArmMeta, self ).initAttributes()
        dep_fn = OpenMaya.MFnDependencyNode()
        dep_fn.setObject( self.mobject )
        
        # Shoulder
        shoudler_str_attr = OpenMaya.MFnTypedAttribute()
        shoulder_obj = shoudler_str_attr.create( 'shoulderMarker', 'shoudlerMarker', OpenMaya.MFnData.kString )
        shoudler_str_attr.setHidden( False )
        dep_fn.addAttribute( shoulder_obj )
        
        # Elbow
        elbow_str_attr = OpenMaya.MFnTypedAttribute()
        elbow_obj = elbow_str_attr.create( 'eldbowMarker', 'elbowMarker', OpenMaya.MFnData.kString )
        elbow_str_attr.setHidden( False )
        dep_fn.addAttribute( elbow_obj )
        
        # Wrist
        wrist_str_attr = OpenMaya.MFnTypedAttribute()
        wrist_obj = wrist_str_attr.create( 'wristMarker', 'wristMarker', OpenMaya.MFnData.kString )
        wrist_str_attr.setHidden( False )
        dep_fn.addAttribute( wrist_obj )
        
    def getAttributes( self ):
        '''
        Returns a dict of attribute name : value.
        '''
        dep_fn = OpenMaya.MFnDependencyNode()
        results = {}
        
        shoulder = getPlugByName( self.mobject, 'shoulderMarker' )
        shoulder_plugs = getPlugConnections( shoulder, True, False )
        if shoulder_plugs:
            dep_fn.setObject( shoulder_plugs[0].node() )
            results[ shoulder.partialName() ] = dep_fn.name()
        else:
            results[ shoulder.partialName() ] = None
        
        elbow = getPlugByName( self.mobject, 'elbowMarker' )
        elbow_plugs = getPlugConnections( elbow, True, False )
        if elbow_plugs:
            dep_fn.setObject( elbow_plugs[0].node() )
            results[ elbow.partialName() ] = dep_fn.name()
        else:
            results[ elbow.partialName() ] = None
        
        wrist = getPlugByName( self.mobject, 'wristMarker' )
        wrist_plugs = getPlugConnections( wrist, True, False )
        if wrist_plugs:
            dep_fn.setObject( wrist_plugs[0].node() )
            results[ wrist.partialName() ] = dep_fn.name()
        else:
            results[ wrist.partialName() ] = None
        
        return results
        
        
        
        
        
rig_node = RigMeta( node_name='rig' )
arm_node = ArmMeta( node_name='sam' )

source_plug = getPlugByName( arm_node.mobject, 'name' )
rig_node.addChildMeta( source_plug )

all_metas = rig_node.getMetaNodes()
for meta in all_metas:
    attrs = meta.getAttributes()
    for key, value in attrs.iteritems():
        print key, value

test_obj = getMObject( 'pSphere1' )

connectNodes( test_obj, 'metanode', arm_node.mobject, 'shoudlerMarker'  )

shoulder_marker_plug = getPlugByName( arm_node.mobject, 'shoudlerMarker' )



















